import os
import pickle
import pandas as pd
from utils import show_time, check_data
from utils_func.env import ROOT_PATH
from utils_func.plots import plot_all
from constant import methods,universes

if __name__ == '__main__':

    # 输入统计项目
    universe_id = input("请选择作图项目(a:苏宁大陆华泰,b:苏宁香港人保,c:苏宁大陆人保,d:联想人保):")# TODO 可根据需求添加universe
    universe = universes[universe_id]['cn_name']
    res_path = os.path.join(ROOT_PATH,f'result\\{universe}')

    lp_raw = pd.read_csv(f'{ROOT_PATH}\data\\2_理赔原始数据_{universe}.csv',encoding='GBK')#_无批次月份
    lp_raw['维修申请时间'] = pd.to_datetime(lp_raw['维修申请时间']).dt.strftime('%Y-%m-%d')
    lp, lp_abnormal = check_data(lp_raw) #TODO save to revise abnormal values

    # 输入风险分摊方法
    method_id = input("请选择风险分摊方法序号(a.时间均摊法 b.时间总和法 c.时间逆总和法):")# TODO 可根据需求添加method
    method = methods[method_id]['cn_name']

    # 结果子目录
    res_subpath = os.path.join(res_path,f'{method}')

    plot_files = {}
    if not os.path.exists(res_subpath):
        print('画图数据不存在,请返回检查是否执行prediction.py')
    else:
        files = os.listdir(res_subpath)
        # latest_time = ''
        plot_files['BAOFEI'] = [f for f in files if f.split('_')[0] == 'BAOFEI']
        plot_files['JINGFEI'] = [f for f in files if f.split('_')[0] == 'JINGFEI']

    fee_dict = {}
    for type in ['BAOFEI','JINGFEI']:
        df = pd.DataFrame()
        for f in plot_files[type]:
            time_range = f.split('_')[-1].split('.')[0]
            df_i = pd.read_csv(os.path.join(res_subpath,f),index_col=0)
            df_i.columns = [f'{type}{time_range}']
            df = pd.concat([df,df_i],axis=1)
        df.fillna(0,inplace=True)
        df_sum = df.sum(axis=1)
        # latest_time = max(latest_time, time_range.split('-')[1])
        fee_dict[type] = pd.DataFrame(df_sum)
    

    show_time('画图开始')
    plot_all(fee_dict['BAOFEI'],fee_dict['JINGFEI'],lp,res_subpath,method)
    show_time('画图完成')